from collections import Counter
from math import sqrt
import numpy as np

class KNNClassifier():
    def __init__(self,k):
        assert k>=1,"k must larger than 0"
        self._k=k
        self._X_train=None
        self._Y_train=None

    def fit(self,X_train,Y_train):
        self._X_train=X_train
        self._Y_train=Y_train
        return self

    def predict(self,X_predict,p=2,mode=0):
        '''
        this function can also achieve regress and classify
        :param X_predict:
        :param p:Minkowski distance parameter
        :param mode:0-classify;1-regress
        :return: regress or predict result
        '''
        return [self._predixt(x,p,mode) for x in X_predict]

    def _predixt(self,x,p,mode):
        # distance=[np.sum((x_train-x) ** int(p))**(1.0/p) for x_train in self._X_tc
        distance = [np.sum((x_train - x) ** int(p)) for x_train in self._X_train]
        nearest=np.argsort(distance)
        topK_Y=[self._Y_train[i] for i in nearest[:self._k]]
        if(mode==0):
            votes=Counter(topK_Y)
            predict_y=votes.most_common(1)[0][0]
        else:
            predict_y=np.mean(topK_Y)
        return predict_y
            # most_common(n) 统计出现最多次数的n个元素

from sklearn import datasets
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

if __name__ == "__main__":
    iris = datasets.load_iris()
    # classify use iris_data to classify the target
    # X = iris.data
    # Y = iris.target
    #regress use iris_data first 3 column to predict the fourth column
    X=iris.data[:,:-1]
    Y=iris.data[:,-1]
    X_train,X_test,Y_train,Y_test=train_test_split(X,Y,test_size=0.2)

    knn=KNNClassifier(k=5)
    knn.fit(X_train,Y_train)
    y_predict=knn.predict(X_test)
    print(range(Y_test.size))
    fig = plt.figure()
    fig1 = fig.add_subplot(221)
    fig1.scatter(X_test[:,0],y_predict)
    fig1.scatter(X_test[:,0],Y_test)
    fig2 = fig.add_subplot(222)
    fig2.scatter(X_test[:,1],y_predict)
    fig2.scatter(X_test[:,1],Y_test)
    fig3 = fig.add_subplot(223)
    fig3.scatter(X_test[:,2],y_predict)
    fig3.scatter(X_test[:,2],Y_test)
    plt.show()

    fig = plt.figure()
    ax = fig.gca(projection='3d')
    ax.scatter(X_test[:,0], X_test[:,1], y_predict, label='parametric curve')
    ax.scatter(X_test[:,0], X_test[:,1], Y_test, label='parametric curve')
    plt.show()

    print(y_predict)
    print(list(Y_test))